import sys
import torch
import random

# 尝试导入 Pygame，如果失败则打印提示
try:
    import pygame
except ImportError:
    print("Error: Pygame is not installed. Please install it using 'pip install pygame'")
    # 提供一个假的 pygame 模块，以便其他代码可以继续执行但会跳过图形部分
    class MockPygame:
        def init(self): pass
        def display(self):
            class MockDisplay:
                def set_mode(self, size): return None
                def set_caption(self, caption): pass
                def flip(self): pass
                def get_active(self): return None
            return MockDisplay()
        def event(self):
            class MockEvent:
                def get(self): return []
                QUIT = -1  # Dummy value
                MOUSEBUTTONDOWN = -2  # Dummy value
            return MockEvent()
        def time(self):
            class MockTime:
                def Clock(self):
                    class MockClock:
                        def tick(self, fps): pass
                    return MockClock()
            return MockTime()
        def font(self):
            class MockFont:
                def Font(self, name, size):
                    class MockRender:
                        def render(self, text, antialias, color): return MockRender()
                        def get_rect(self, center): return MockRender()
                        def blit(self, surface, rect): pass
                    return MockRender()
            return MockFont()
        def quit(self): pass
        K_ESCAPE = -3  # Dummy value
        def draw(self):
            class MockDraw:
                def line(self, s, c, p1, p2, w): pass
                def circle(self, s, c, p, r, w=0): pass
            return MockDraw()

    pygame = MockPygame()


from rainforeLearn.gomoku.v1.dqn.configs.config import DQNConfig
from rainforeLearn.gomoku.v1.dqn.models.q_network import QNetwork, DuelingQNetwork, ConvolutionalDuelingQNetwork
from rainforeLearn.gomoku.v1.dqn.agents.dqn_agent import DQNAgent
from rainforeLearn.gomoku.v1.gomoku_env import GomokuEnv


class GomokuGame:
    """
    五子棋GUI游戏类，允许人类玩家与DQN智能体对弈。
    """

    def __init__(self, model_path=None, board_size=15, ai_player_num=-1):
        # 检查 Pygame 是否可用
        if 'pygame' not in sys.modules or pygame.display.get_active() is None:
            print("Pygame is not initialized or display not available. GUI will not start.")
            self.gui_enabled = False
            return

        self.gui_enabled = True
        try:
            print("GomokuGame: Initializing Pygame...")
            pygame.init()
            print("GomokuGame: Pygame initialized.")
        except Exception as e:
            print(f"Error initializing Pygame: {e}. GUI will not start.")
            self.gui_enabled = False
            return

        self.board_size = board_size
        self.cell_size = 40
        self.padding = 20
        self.screen_width = self.board_size * self.cell_size + self.padding * 2
        self.screen_height = self.board_size * self.cell_size + self.padding * 2 + 60

        try:
            print(f"GomokuGame: Creating screen with dimensions {self.screen_width}x{self.screen_height}...")
            self.screen = pygame.display.set_mode((self.screen_width, self.screen_height))
            pygame.display.set_caption("五子棋 (人 vs AI)")
            print("GomokuGame: Screen created successfully.")
        except Exception as e:
            print(f"Error creating Pygame screen: {e}. GUI will not start.")
            self.gui_enabled = False
            pygame.quit()
            return

        self.font = pygame.font.Font(None, 36)

        self.env = GomokuEnv(self.board_size)
        self.state = self.env.reset()
        self.game_over = False
        self.winner = None

        # 颜色定义
        self.BOARD_COLOR = (240, 200, 140)
        self.LINE_COLOR = (0, 0, 0)
        self.BLACK_STONE_COLOR = (0, 0, 0)
        self.WHITE_STONE_COLOR = (255, 255, 255)

        # AI设置
        self.ai_agent = None
        self.ai_enabled = False
        if model_path:
            try:
                config = DQNConfig()
                # 🔥 关键修复：硬编码与训练时匹配的配置
                config.state_dim = self.env.observation_space_shape[0]
                config.action_dim = self.env.action_space_n
                config.hidden_dim = 512  # 🔥 从错误信息推断出的值
                config.double_dqn = True
                config.dueling_dqn = True
                config.prioritized_replay = True
                config.use_residual_connections = True  # 🔥 关键参数
                config.dropout_rate = 0.15  # 🔥 训练时的默认值
                config.use_noisy_networks = False

                # 🔥 选择正确的网络架构（与训练时保持一致）
                if config.use_conv_network:
                    network_class = ConvolutionalDuelingQNetwork
                    network_factory = lambda state_dim, action_dim, hidden_dim: \
                        network_class(board_size, action_dim, hidden_dim,
                                      config.dropout_rate, config.use_noisy_networks)
                else:
                    # 使用全连接网络
                    if config.dueling_dqn:
                        network_class = DuelingQNetwork
                    else:
                        network_class = QNetwork

                    network_factory = lambda state_dim, action_dim, hidden_dim: \
                        network_class(state_dim, action_dim, hidden_dim,
                                      config.use_residual_connections,
                                      config.dropout_rate,
                                      config.use_noisy_networks)

                # 创建智能体
                self.ai_agent = DQNAgent(config, network_factory)
                self.ai_agent.load(model_path)
                self.ai_agent.policy_net.eval()
                self.ai_enabled = True
                print(f"✅ AI模型 '{model_path}' 加载成功!")
                print(f"📊 使用配置: hidden_dim={config.hidden_dim}, "
                      f"residual={config.use_residual_connections}, "
                      f"dueling={config.dueling_dqn}")

            except Exception as e:
                print(f"❌ 加载AI模型失败: {e}")
                print("💡 提示: 请检查模型文件是否存在且配置匹配")
                self.ai_enabled = False
        else:
            print("未提供AI模型路径，将进行双人对弈。")

        self.human_player = 1
        self.ai_player_num = ai_player_num
        self.current_turn_player = 1

    def render(self, screen=None, cell_size=None, padding=None):
        """用于Pygame GUI绘制棋盘和棋子"""
        if 'pygame' not in sys.modules or screen is None or cell_size is None or padding is None:
            return

        import pygame

        # 绘制棋盘格线
        for i in range(self.board_size):
            pygame.draw.line(screen, (0, 0, 0),
                             (padding + i * cell_size + cell_size / 2, padding + cell_size / 2),
                             (padding + i * cell_size + cell_size / 2,
                              padding + self.board_size * cell_size - cell_size / 2), 2)
            pygame.draw.line(screen, (0, 0, 0),
                             (padding + cell_size / 2, padding + i * cell_size + cell_size / 2),
                             (padding + self.board_size * cell_size - cell_size / 2,
                              padding + i * cell_size + cell_size / 2), 2)

        # 绘制棋子
        for r in range(self.board_size):
            for c in range(self.board_size):
                center_x = padding + c * cell_size + cell_size / 2
                center_y = padding + r * cell_size + cell_size / 2
                stone_radius = int(cell_size / 2.5)

                player_on_board = self.env.board[r, c]
                if player_on_board == 1:
                    pygame.draw.circle(screen, (0, 0, 0), (int(center_x), int(center_y)), stone_radius)
                elif player_on_board == -1:
                    pygame.draw.circle(screen, (255, 255, 255), (int(center_x), int(center_y)), stone_radius)
                    pygame.draw.circle(screen, (0, 0, 0), (int(center_x), int(center_y)), stone_radius, 1)
    def draw_board(self):
        """绘制棋盘、棋子和游戏状态信息"""
        if not self.gui_enabled:
            return

        self.screen.fill(self.BOARD_COLOR)
        self.render(self.screen, self.cell_size, self.padding)

        # 显示游戏状态
        message = ""
        if self.game_over:
            if self.winner == self.human_player:
                message = "You Win!"
            elif self.winner == self.ai_player_num:
                message = "AI Wins!"
            else:
                message = "Draw!"
            message += " (Click to restart)"
        else:
            player_name = "Your" if self.current_turn_player == self.human_player else "AI's"
            color_name = "Black" if self.current_turn_player == 1 else "White"
            message = f"{player_name} turn ({color_name})"

        text_surface = self.font.render(message, True, (0, 0, 0))
        text_rect = text_surface.get_rect(center=(self.screen_width // 2, self.screen_height - 30))
        self.screen.blit(text_surface, text_rect)

        pygame.display.flip()

    def handle_click(self, pos):
        """处理鼠标点击事件"""
        if self.game_over:
            return None

        x, y = pos
        x -= self.padding
        y -= self.padding

        col = int(x / self.cell_size)
        row = int(y / self.cell_size)

        if not (0 <= row < self.board_size and 0 <= col < self.board_size):
            return None

        action = row * self.board_size + col
        return action

    def play(self):
        """游戏主循环"""
        if not self.gui_enabled:
            print("GUI is not enabled. Cannot start game loop.")
            return

        print("GomokuGame: Entering game loop...")
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                    if self.game_over:
                        self.state = self.env.reset()
                        self.game_over = False
                        self.winner = None
                        self.current_turn_player = 1
                        if self.ai_enabled and self.current_turn_player == self.ai_player_num:
                            self.ai_turn()
                        continue

                    if self.current_turn_player == self.human_player:
                        action = self.handle_click(event.pos)
                        if action is not None:
                            self.make_move(action, self.human_player)
                            if not self.game_over and self.ai_enabled and self.current_turn_player == self.ai_player_num:
                                self.ai_turn()

            self.draw_board()
            pygame.time.Clock().tick(30)

        print("GomokuGame: Exiting game loop. Quitting Pygame...")
        pygame.quit()
        sys.exit()

    def make_move(self, action, player):
        """执行落子动作"""
        if self.game_over:
            return

        self.env.current_player = player

        next_state, reward, done, info = self.env.step(action)

        self.state = next_state
        self.game_over = done
        self.winner = self.env.winner
        self.current_turn_player *= -1

        if 'invalid_move' in info and info['invalid_move']:
            print(f"玩家 {player} 进行了非法落子。游戏结束。")
            self.game_over = True

        self.draw_board()

    def ai_turn(self):
        """🔥 优化的AI玩家回合逻辑 - 确保不会选择非法动作"""
        if not self.ai_enabled or self.game_over:
            return

        print("AI思考中...")

        legal_actions = self.env.get_legal_actions()

        if not legal_actions:
            print("没有合法位置可以落子！")
            return

        state = self.env._get_state()

        with torch.no_grad():
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.ai_agent.device)
            q_values = self.ai_agent.policy_net(state_tensor).squeeze(0)

        # 🔥 动作掩码：确保只选择合法动作
        masked_q_values = q_values.clone()
        all_actions = set(range(self.env.action_space_n))
        illegal_actions = all_actions - set(legal_actions)

        for action in illegal_actions:
            masked_q_values[action] = float('-inf')

        ai_action = masked_q_values.argmax().item()

        # 双重检查：确保选择的动作是合法的
        if ai_action not in legal_actions:
            print(f"警告：AI选择了非法动作 {ai_action}，改为随机选择合法动作")
            ai_action = random.choice(legal_actions)

        self.make_move(ai_action, self.ai_player_num)
        print(f"AI落子于: ({ai_action // self.board_size}, {ai_action % self.board_size})")